import torch
from torch import nn
from src.model.space_model.dgraph import DGraph
from src.model.space_model.agraph import AGraph
from src.model.space_model.gcn import Gcn


class SpaceBlock(nn.Module):
    def __init__(self, dims, seq_len, num_nodes, num_layer, static_graph, single_point=True):
        super(SpaceBlock, self).__init__()
        self.num_nodes = num_nodes
        self.seq_len = seq_len
        self.num_layer = num_layer
        self.input_size, self.hidden_size, self.output_size = dims
        self.static_graph = static_graph
        self.single_point = single_point
        if not self.static_graph:
            self.dgraph = DGraph((self.hidden_size * self.seq_len // 2, self.hidden_size, self.output_size), self.num_nodes)
            self.agraph = AGraph(dims, self.seq_len, self.num_nodes)
        self.gcn = Gcn(dims)
        self.weights = nn.Parameter(torch.tensor([0.33, 0.33, 0.33]))

    def forward(self, x, y, adj, is_training):
        # 若不是单点预测，adj尺寸为(B, N, N), x的尺寸是(B, P, N, H)
        if not self.single_point:
            adj = adj.unsqueeze(1).repeat(1, y.size(1), 1, 1)
        
        if not self.static_graph:
            # x.shape = (B, T, N, H)
            adj_a = self.agraph(x)
            x = x.permute(0, 2, 1, 3).contiguous()
            dx = x.reshape(x.shape[0], x.shape[1], -1)
            adj_d = self.dgraph(dx, adj)
            
            
            weights = torch.softmax(self.weights, dim=0)
            adj_g = adj_a * weights[0]  + adj_d * weights[1] + adj * weights[2]

        # gcn的输出out是(B, N, H)
        out = self.gcn(y, adj_g)
        if is_training:
            return out, adj_g
        else:
            print(weights)
            return out, adj_g, adj_a, adj_d, adj